// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

{% macro per_exec_ws(folder) -%}
  "workspace/exec_${env.EXECUTOR_NUMBER}/{{ folder }}"
{%- endmacro -%}

{% macro junit_to_s3() %}
sh(
            script: """
              set -eux
              aws s3 cp --no-progress build/pytest-results s3://${s3_prefix}/pytest-results --recursive
            """,
            label: 'Upload JUnits to S3',
          )
{% endmacro %}

{% macro sharded_test_step(name, num_shards, node, ws, docker_image, platform, test_method_names) %}

{% for shard_index in range(1, num_shards + 1) %}
{% set method_name = "shard_run_" + name.replace(":", "").replace(" ", "-").replace("-", "_") + "_" + shard_index|string + "_of_" + num_shards|string %}
def {{ method_name }}() {
  if (!skip_ci && is_docs_only_build != 1) {
    node('{{ node }}') {
      ws({{ per_exec_ws(ws) }}) {
        try {
          docker_init({{ docker_image }})
          init_git()
          timeout(time: max_time, unit: 'MINUTES') {
            withEnv([
              'PLATFORM={{ platform }}',
              'TVM_NUM_SHARDS={{ num_shards }}',
              'TVM_SHARD_INDEX={{ shard_index - 1 }}'], {
              {{ caller(shard_index, num_shards) | trim | indent(width=12) }}
            })
          }
        } finally {
          {{ junit_to_s3() }}
          junit 'build/pytest-results/*.xml'
        }
      }
    }
  } else {
    Utils.markStageSkippedForConditional('{{ name }} {{ shard_index }} of {{ num_shards }}')
  }
}
{% set _ = test_method_names.append((name + " " + shard_index|string + " of " + num_shards|string, method_name)) %}

{% endfor %}
{% endmacro %}

{% macro sharded_lint_step(name, num_shards, docker_image, node, ws) %}
{% for shard_index in range(1, num_shards + 1) %}
  '{{ name }} {{ shard_index }} of {{ num_shards }}': {
    node('{{ node }}') {
      ws({{ per_exec_ws(ws) }}) {
        docker_init({{ docker_image }})
        init_git()
        timeout(time: max_time, unit: 'MINUTES') {
          withEnv([
            'TVM_NUM_SHARDS={{ num_shards }}',
            'TVM_SHARD_INDEX={{ shard_index - 1 }}'], {
            {{ caller() | trim | indent(width=6) }}
          })
        }
      }
    }
  },
{% endfor %}
{% endmacro %}


{% macro test_step(name, node, ws, docker_image, platform) %}
  '{{ name }}': {
    if (!skip_ci && is_docs_only_build != 1) {
      node('{{ node }}') {
        ws({{ per_exec_ws(ws) }}) {
          timeout(time: max_time, unit: 'MINUTES') {
            try {
              docker_init({{ docker_image }})
              init_git()
              withEnv(['PLATFORM={{ platform }}'], {
                {{ caller() | indent(width=12) | trim }}
              })
            } finally {
              {{ junit_to_s3() | indent(width=4) }}
              junit 'build/pytest-results/*.xml'
            }
          }
        }
      }
    } else {
      Utils.markStageSkippedForConditional('{{ name }}')
    }
  },
{% endmacro %}

{% macro upload_artifacts(tag, filenames, folders=None) %}
sh(
            script: """
              set -eux
              {% for filename in filenames %}
              md5sum {{ filename }}
              aws s3 cp --no-progress {{ filename }} s3://${s3_prefix}/{{ tag }}/{{ filename }}
              {% endfor %}
              {% for folder in (folders or []) %}
              aws s3 cp --no-progress {{ folder }} s3://${s3_prefix}/{{ tag }}/{{ folder }} --recursive
              {% endfor %}
            """,
            label: 'Upload artifacts to S3',
          )
{% endmacro %}

{% macro download_artifacts(tag, filenames, folders=None) %}
sh(
            script: """
              set -eux
              {% for filename in filenames %}
              aws s3 cp --no-progress s3://${s3_prefix}/{{ tag }}/{{ filename }} {{ filename }}
              md5sum {{ filename }}
              {% endfor %}
              {% for folder in (folders or []) %}
              aws s3 cp --no-progress s3://${s3_prefix}/{{ tag }}/{{ folder }} {{ folder }} --recursive
              {% endfor %}
            """,
            label: 'Download artifacts from S3',
          )
{% endmacro %}
